package com.liuhuan.utils;
import java.io.BufferedReader;
import java.io.DataOutputStream;
import java.io.InputStreamReader;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.Scanner;
import org.json.JSONArray;
import org.json.JSONObject;

public class Main {
    private static final String USER_AGENT = "Java-HttpURLConnection/1.0";
    private static final String API_KEY = "sk-6e52620471f54ff0ba86b264c02eb44b";
    public static void main(String[] args) throws Exception {
        // 定义messages
        JSONArray messages = new JSONArray();
        // 定义系统消息system_message并添加到messages中
        JSONObject system_message = new JSONObject();
        system_message.put("role","system");
        system_message.put("content","你是刘欢研发的代码生成助手，能生成crud基础代码.");
        messages.put(system_message);
        // 您可以在此修改对话轮数，当前为3轮对话
        for (int i = 0; i < 5; i++) {
            Scanner scanner = new Scanner(System.in);
            System.out.print("请输入：");
            String UserInput = scanner.next();
            // 定义用户消息user_message并添加到messages中
            JSONObject user_message = new JSONObject();
            user_message.put("role","user");
            user_message.put("content",UserInput);
            messages.put(user_message);
            // 定义传入的body对象
            JSONObject body = new JSONObject();
            // 定义body中的input对象
            JSONObject input_message = new JSONObject();
            input_message.put("messages",messages);
            // 定义body中的parameters对象
            JSONObject parameters_message = new JSONObject();
            parameters_message.put("result_format","message");
            // 将model、input与parameters传入body中
            body.put("model","qwen-turbo");
            body.put("input",input_message);
            body.put("parameters",parameters_message);

            // 进行模型的调用
            JSONObject assistant_output = get_response(body);
            // 取出content并打印出来
            JSONObject output = assistant_output.getJSONObject("output");
            JSONArray choices = output.getJSONArray("choices");
            JSONObject firstChoice = choices.getJSONObject(0);
            JSONObject assistant_message = firstChoice.getJSONObject("message");
            String content = assistant_message.getString("content");
            System.out.println("用户输入："+UserInput);
            System.out.println("模型输出："+content+'\n');
            // 将assistant_message添加到messages中
            messages.put(assistant_message);
        }
    }
    // 封装模型响应函数
    public static JSONObject get_response(JSONObject body) throws Exception {
        String urlStr = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation";
        String apiKey = API_KEY;
        URL url = new URL(urlStr);
        HttpURLConnection connection = (HttpURLConnection) url.openConnection();

        connection.setRequestMethod("POST");
        connection.setRequestProperty("Content-Type", "application/json");
        connection.setRequestProperty("Authorization", "Bearer " + apiKey);
        connection.setDoOutput(true);

        String jsonInputString = body.toString();
        try (DataOutputStream wr = new DataOutputStream(connection.getOutputStream())) {
            wr.write(jsonInputString.getBytes(StandardCharsets.UTF_8));
            wr.flush();
        }

        StringBuilder response = new StringBuilder();
        try (BufferedReader in = new BufferedReader(
                new InputStreamReader(connection.getInputStream()))) {
            String inputLine;
            while ((inputLine = in.readLine()) != null) {
                response.append(inputLine);
            }
        }
        JSONObject response_json = new JSONObject(response.toString());
        connection.disconnect();
        return response_json;
    }
}